package cn.myeasyai.core.temple;

import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.dromara.easyai.yolo.TypeBody;
import org.dromara.easyai.yolo.YoloBody;
import org.dromara.easyai.yolo.YoloConfig;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class Temple {

    public MyTemples readALlTemple(List<AllTemple> allTemples, YoloConfig yoloConfig) throws Exception {
        MyTemples myTemples = new MyTemples();
        TypeBody trTypeBody = new TypeBody(yoloConfig, yoloConfig.getWindowWidth(), yoloConfig.getWindowHeight());
        trTypeBody.setTypeID(2);
        TypeBody faceTypeBody = new TypeBody(yoloConfig, yoloConfig.getWindowWidth(), yoloConfig.getWindowHeight());
        trTypeBody.setTypeID(1);
        List<TempleMessage> templeMessages = new ArrayList<>();
        myTemples.setTempleMessages(templeMessages);
        myTemples.setTrTypeBody(trTypeBody);
        for (AllTemple temple : allTemples) {
            String templeUrl = temple.getTempleUrl();
            String pictureUrl = temple.getPictureUrl();
            ////
            File folder = new File(templeUrl);
            File[] files = folder.listFiles();
            if (files != null) {
                for (File file : files) {
                    if (file.getName().contains("xml")) {
                        TempleMessage templeMessage = readTemple(file, pictureUrl, trTypeBody, faceTypeBody);
                        templeMessages.add(templeMessage);
                    }
                }
            }
        }
        myTemples.setFaceTypeBody(faceTypeBody);
        return myTemples;
    }

    private TempleMessage readTemple(File file, String url, TypeBody trTypeBody, TypeBody faceTypeBody) throws Exception {
        int all = 0;
        TempleMessage templeMessage = new TempleMessage();
        YoloBody faceYoloBody = new YoloBody();
        YoloBody trYoloBody = new YoloBody();
        templeMessage.setFaceYoloBody(faceYoloBody);
        templeMessage.setTrYoloBody(trYoloBody);
        SAXReader saxReader = new SAXReader();
        Document document = saxReader.read(file);
        Element rootElement = document.getRootElement();
        Element element = rootElement.element("filename");
        List<Element> elements = rootElement.elements("object");
        for (Element myElement : elements) {
            String name = myElement.element("name").getData().toString();
            Element sonElement = myElement.element("bndbox");
            int xMin = Integer.parseInt(sonElement.element("xmin").getData().toString());
            int yMin = Integer.parseInt(sonElement.element("ymin").getData().toString());
            int xMax = Integer.parseInt(sonElement.element("xmax").getData().toString());
            int yMax = Integer.parseInt(sonElement.element("ymax").getData().toString());
            int width = xMax - xMin;
            int height = yMax - yMin;
            if (name.equals("face")) {
                all = all | 1;
                faceYoloBody.setX(xMin);
                faceYoloBody.setY(yMin);
                faceYoloBody.setWidth(width);
                faceYoloBody.setHeight(height);
                faceYoloBody.setTypeID(1);
                faceTypeBody.insertYoloBody(faceYoloBody);
            } else if (name.equals("tr")) {
                all = all | (1 << 1);
                trYoloBody.setX(xMin);
                trYoloBody.setY(yMin);
                trYoloBody.setWidth(width);
                trYoloBody.setHeight(height);
                trYoloBody.setTypeID(2);
                trTypeBody.insertYoloBody(trYoloBody);
            } else {
                throw new Exception("标注名称出现错误:" + element.getData().toString());
            }
        }
        if (all != 3) {
            throw new Exception("标注名称出现遗漏:" + element.getData().toString());
        }
        templeMessage.setFileName(url + "/" + element.getData().toString());
        return templeMessage;
    }
}
